import matplotlib.pyplot as plt
import torch
import torch
import numpy as np
from sklearn.datasets import load_boston
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split

(data, target) = load_boston(return_X_y=True)
# data = boston.data
target = target.reshape(-1, 1)

data = MinMaxScaler().fit_transform(data)
target = MinMaxScaler().fit_transform(target)

c = 7
x = []
y = []
for i in range(len(data) - c):
    x_data = data[i:i + c]
    y_data = target[i + c]
    x.append(x_data)
    y.append(y_data)

x = torch.Tensor(x)
y = torch.Tensor(y)

train_x, test_x, train_y, test_y = train_test_split(x, y, shuffle=False)
print(train_x.shape)

# 定义模型
model = torch.nn.Linear(in_features=7 * 13, out_features=1)

# 定义损失函数和优化器
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters())

# 训练模型
loss_list = []
for epoch in range(2000):
    model.zero_grad()
    outputs = model(train_x.reshape(-1, 7 * 13))
    loss = loss_fn(outputs, train_y)
    loss_list.append(loss.item())
    loss.backward()
    optimizer.step()
    if epoch % 10 == 0:
        print(f'Epoch {epoch + 1}, Loss: {loss.item()}')

# 预测测试集结果
test_outputs = model(test_x.reshape(-1, 7 * 13)).data.numpy()
model.eval()
with torch.no_grad():
    # 绘制结果
    plt.plot(test_y.data.numpy(), c='r')
    plt.plot(test_outputs, c='g')
    plt.show()
